import os
import requests
from langchain_core.tools import Tool


bocha_api_key = os.environ.get('BOCHA_KEY')
# print(bocha_api_key)

def _bocha_web_search_tool(query: str, count: int = 5) -> str:
    """
    使用 bocha web search api 进行联网搜索，返回搜索结果的字符串。
    参数:
    - query: 搜索关键词
    - count: 返回的搜索结果数量
    返回:
    - 搜索结果的字符串形式
    """
    url = 'https://api.bochaai.com/v1/web-search'
    headers = {
         'Authorization': f'Bearer {bocha_api_key}',
         'Content-Type': 'application/json'
    }
    data = {
        "query": query,
        "freshness": "nolimit",  # 搜索的时间范围，例如"oneday","oneweek","onemonth","oneyear","nolimit"
        "summary": True,  # 是否返回长文本摘要总结
        "count": count
    }

    response = requests.post(url, headers=headers, json=data)
    if response.status_code == 200:
        # 返回给大模型的格式化的搜索结果文本
        # 可以自己对博查的搜索结果进行自定义处理
        data = response.json()
        web_pages = data['data']['webPages']['value']
        return "\n\n".join([page['summary']for page in web_pages])
    else:
        raise Exception(f"API 请求失败，状态码: {response.status_code}, 错误信息: {response.text}")



# 创建 langchain 工具
bocha_tool = Tool(
    name="bochawebsearch",
    func=_bocha_web_search_tool,
    description="使用 bocha web search api 进行网络搜索,  当需要查询当前事件时可以使用这个工具。"
)